from graphviz import Digraph


def draw_rnn():
    dot = Digraph()

    # 定义节点
    dot.node('x', 'Input (x)')
    dot.node('h0', 'Initial Hidden State (h0)')
    dot.node('rnn', 'RNN Cell')
    dot.node('h', 'Hidden State (h)')
    dot.node('out', 'RNN Output (out)')
    dot.node('fc', 'Fully Connected Layer')
    dot.node('y', 'Output (y)')

    # 定义边
    dot.edges(['xh0', 'hx', 'hout'])
    dot.edge('out', 'fc')
    dot.edge('fc', 'y')

    # 增加输入节点和隐藏状态的连接
    dot.edge('x', 'rnn', label='x(t)')
    dot.edge('h0', 'rnn', label='h(t-1)')
    dot.edge('rnn', 'h', label='h(t)')
    dot.edge('rnn', 'out', label='out(t)')

    # 设置图的方向从左到右
    dot.attr(rankdir='LR')

    # 绘制图形
    dot.render('simple_rnn_graph', format='png', view=True)


draw_rnn()
